import pickle
import cv2
import numpy as np


class Cifar10Preprocess:
    """
    Dataset loader for CIFAR-10.

    Args:
        path(str): dataset path
        dsize(tuple, integer): output shape of resize operation
        rescale(float): rescale factor of rescale operation
        shift(float): shift factor of rescale operation
        mean(tuple): mean of normalize
        std(tuple): variance of normalize
    """
    def __init__(self, path, dsize=224, rescale=1.0/255.0, shift=0.0,
                 mean=None, std=None, batch_size=32):
        # parameters for data preprocess
        if isinstance(dsize, int):
            self.dsize = (dsize, dsize)
        elif isinstance(dsize, tuple):
            self.dsize = dsize
        else:
            raise TypeError(f"Unsupported data type of `dsize`: {type(dsize)}")
        self.rescale = rescale
        self.shift = shift
        self.mean = mean
        self.std = std
        self.batch_size = batch_size
        # load dataset
        lines = np.fromfile(path, dtype=np.uint8).reshape(10000, 3073)
        # do preprocess and batch
        self.samples = []
        self.labels = []
        batch = []
        label_batch = []
        for line in lines:
            if len(batch) < self.batch_size:
                label_batch.append(line[0])
                batch.append(self.__preprocess(line[1:]))
            else:
                self.samples.append(batch)
                batch = [self.__preprocess(line[1:])]
                self.labels.append(label_batch)
                label_batch = [line[0]]

        if len(batch) == self.batch_size:
            self.samples.append(batch)
            self.labels.append(label_batch)
        self.num_samples = len(self.labels)

    def __preprocess(self, line):
        image = self.__decode(line)
        image = self.__resize(image)
        # TODO: rescale and normalize may cause difference
        image = self.__rescale(image)
        image = self.__normalize(image)
        return image

    def __decode(self, line):
        """Decode image."""
        image = np.reshape(line, newshape=(3, 1024)).T
        image = np.reshape(image, newshape=(32, 32, 3))
        # TODO: whether it need to convert BGR to RGB
        # image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        return image

    def __resize(self, image):
        """Resize image."""
        # TODO: whether it is INTER_LINEAR
        image = cv2.resize(image, dsize=self.dsize, interpolation=cv2.INTER_LINEAR)
        # cv2.resize(image, dsize=(224, 224), interpolation=cv2.INTER_CUBIC)
        return image

    def __rescale(self, image):
        """Rescale array."""
        return image * self.rescale + self.shift

    def __normalize(self, image):
        """Normalize array"""
        return (image - self.mean) / self.std

    def __getitem__(self, item):
        if item < self.num_samples:
            image = self.samples[item]
            label = self.labels[item]
        else:
            raise StopIteration
        return image, label


if __name__ == "__main__":
    cifar_preprocess = Cifar10Preprocess("./cifar-10-batches-bin/test_batch.bin",
                                         dsize=224, rescale=1.0/255.0, shift=0.0,
                                         mean=(0.4914, 0.4822, 0.4465),
                                         std=(0.2023, 0.1994, 0.2010),
                                         batch_size=32)
    for idx, sample in enumerate(cifar_preprocess):
        print(type(sample[0]))
        print(len(sample[0]), sample[0][0].shape)
        print(type(sample[1]))
        break
